{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7fcc021a",
   "metadata": {},
   "source": [
    "<img src=\"http://developer.download.nvidia.com/notebooks/dlsw-notebooks/tensorrt_torchtrt_efficientnet/nvidia_logo.png\" width=\"90px\">\n",
    "\n",
    "# Pyspark TensorFlow Inference\n",
    "\n",
    "### Classification using Keras Preprocessing Layers\n",
    "\n",
    "In this notebook, we demonstrate distributed inference using Keras preprocessing layers to classify structured data.  \n",
    "From: https://www.tensorflow.org/tutorials/structured_data/preprocessing_layers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35203476",
   "metadata": {},
   "source": [
    "Note that cuFFT/cuDNN/cuBLAS registration errors are expected (as of `tf=2.17.0`) and will not affect behavior, as noted in [this issue.](https://github.com/tensorflow/tensorflow/issues/62075)  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "01162f42-0637-4dfe-8d7d-b577e4ffd017",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-02-04 13:59:29.670948: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "2025-02-04 13:59:29.679838: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:485] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "2025-02-04 13:59:29.689914: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:8454] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "2025-02-04 13:59:29.692851: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1452] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "2025-02-04 13:59:29.700499: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 AVX_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2025-02-04 13:59:30.139239: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import shutil\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "\n",
    "from tensorflow.keras import layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0d586fb8",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.mkdir('models') if not os.path.exists('models') else None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fa3e1b7-58cd-45f9-9fee-85f25a31c3c6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.17.0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: All log messages before absl::InitializeLog() is called are written to STDERR\n",
      "I0000 00:00:1738706370.524690 3686377 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "I0000 00:00:1738706370.550329 3686377 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "I0000 00:00:1738706370.553239 3686377 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n"
     ]
    }
   ],
   "source": [
    "print(tf.__version__)\n",
    "\n",
    "# Enable GPU memory growth\n",
    "gpus = tf.config.experimental.list_physical_devices('GPU')\n",
    "if gpus:\n",
    "    try:\n",
    "        for gpu in gpus:\n",
    "            tf.config.experimental.set_memory_growth(gpu, True)\n",
    "    except RuntimeError as e:\n",
    "        print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2402b9a",
   "metadata": {},
   "source": [
    "#### Download dataset\n",
    "\n",
    "Download the PetFinder dataset from Kaggle, which where each row describes a pet and the goal is to predict adoption speed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9326b072-a53c-40c4-a6cb-bd4d3d644d03",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading data from http://storage.googleapis.com/download.tensorflow.org/data/petfinder-mini.zip\n",
      "\u001b[1m1668792/1668792\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m0s\u001b[0m 0us/step\n"
     ]
    }
   ],
   "source": [
    "import pathlib\n",
    "import os\n",
    "dataset_url = 'http://storage.googleapis.com/download.tensorflow.org/data/petfinder-mini.zip'\n",
    "\n",
    "data_dir = tf.keras.utils.get_file('petfinder_mini.zip', dataset_url, extract=True, cache_dir='.')\n",
    "data_dir = pathlib.Path(data_dir)\n",
    "try:\n",
    "    # pet-finder-mini might be under a parent a directory petfinder_mini_extracted. Check if this is the case:\n",
    "    dataset = os.path.join(os.path.dirname(data_dir), 'petfinder_mini_extracted/petfinder-mini/petfinder-mini.csv')\n",
    "    dataframe = pd.read_csv(dataset)\n",
    "except:\n",
    "    dataset = os.path.join(os.path.dirname(data_dir), 'petfinder-mini/petfinder-mini.csv')\n",
    "    dataframe = pd.read_csv(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e98480ef-d13d-44c0-a227-e9a22f9bf2b0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Type</th>\n",
       "      <th>Age</th>\n",
       "      <th>Breed1</th>\n",
       "      <th>Gender</th>\n",
       "      <th>Color1</th>\n",
       "      <th>Color2</th>\n",
       "      <th>MaturitySize</th>\n",
       "      <th>FurLength</th>\n",
       "      <th>Vaccinated</th>\n",
       "      <th>Sterilized</th>\n",
       "      <th>Health</th>\n",
       "      <th>Fee</th>\n",
       "      <th>Description</th>\n",
       "      <th>PhotoAmt</th>\n",
       "      <th>AdoptionSpeed</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Cat</td>\n",
       "      <td>3</td>\n",
       "      <td>Tabby</td>\n",
       "      <td>Male</td>\n",
       "      <td>Black</td>\n",
       "      <td>White</td>\n",
       "      <td>Small</td>\n",
       "      <td>Short</td>\n",
       "      <td>No</td>\n",
       "      <td>No</td>\n",
       "      <td>Healthy</td>\n",
       "      <td>100</td>\n",
       "      <td>Nibble is a 3+ month old ball of cuteness. He ...</td>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Cat</td>\n",
       "      <td>1</td>\n",
       "      <td>Domestic Medium Hair</td>\n",
       "      <td>Male</td>\n",
       "      <td>Black</td>\n",
       "      <td>Brown</td>\n",
       "      <td>Medium</td>\n",
       "      <td>Medium</td>\n",
       "      <td>Not Sure</td>\n",
       "      <td>Not Sure</td>\n",
       "      <td>Healthy</td>\n",
       "      <td>0</td>\n",
       "      <td>I just found it alone yesterday near my apartm...</td>\n",
       "      <td>2</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Dog</td>\n",
       "      <td>1</td>\n",
       "      <td>Mixed Breed</td>\n",
       "      <td>Male</td>\n",
       "      <td>Brown</td>\n",
       "      <td>White</td>\n",
       "      <td>Medium</td>\n",
       "      <td>Medium</td>\n",
       "      <td>Yes</td>\n",
       "      <td>No</td>\n",
       "      <td>Healthy</td>\n",
       "      <td>0</td>\n",
       "      <td>Their pregnant mother was dumped by her irresp...</td>\n",
       "      <td>7</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Dog</td>\n",
       "      <td>4</td>\n",
       "      <td>Mixed Breed</td>\n",
       "      <td>Female</td>\n",
       "      <td>Black</td>\n",
       "      <td>Brown</td>\n",
       "      <td>Medium</td>\n",
       "      <td>Short</td>\n",
       "      <td>Yes</td>\n",
       "      <td>No</td>\n",
       "      <td>Healthy</td>\n",
       "      <td>150</td>\n",
       "      <td>Good guard dog, very alert, active, obedience ...</td>\n",
       "      <td>8</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Dog</td>\n",
       "      <td>1</td>\n",
       "      <td>Mixed Breed</td>\n",
       "      <td>Male</td>\n",
       "      <td>Black</td>\n",
       "      <td>No Color</td>\n",
       "      <td>Medium</td>\n",
       "      <td>Short</td>\n",
       "      <td>No</td>\n",
       "      <td>No</td>\n",
       "      <td>Healthy</td>\n",
       "      <td>0</td>\n",
       "      <td>This handsome yet cute boy is up for adoption....</td>\n",
       "      <td>3</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  Type  Age                Breed1  Gender Color1    Color2 MaturitySize  \\\n",
       "0  Cat    3                 Tabby    Male  Black     White        Small   \n",
       "1  Cat    1  Domestic Medium Hair    Male  Black     Brown       Medium   \n",
       "2  Dog    1           Mixed Breed    Male  Brown     White       Medium   \n",
       "3  Dog    4           Mixed Breed  Female  Black     Brown       Medium   \n",
       "4  Dog    1           Mixed Breed    Male  Black  No Color       Medium   \n",
       "\n",
       "  FurLength Vaccinated Sterilized   Health  Fee  \\\n",
       "0     Short         No         No  Healthy  100   \n",
       "1    Medium   Not Sure   Not Sure  Healthy    0   \n",
       "2    Medium        Yes         No  Healthy    0   \n",
       "3     Short        Yes         No  Healthy  150   \n",
       "4     Short         No         No  Healthy    0   \n",
       "\n",
       "                                         Description  PhotoAmt  AdoptionSpeed  \n",
       "0  Nibble is a 3+ month old ball of cuteness. He ...         1              2  \n",
       "1  I just found it alone yesterday near my apartm...         2              0  \n",
       "2  Their pregnant mother was dumped by her irresp...         7              3  \n",
       "3  Good guard dog, very alert, active, obedience ...         8              2  \n",
       "4  This handsome yet cute boy is up for adoption....         3              2  "
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataframe.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27d844f1",
   "metadata": {},
   "source": [
    "### Prepare dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e8efce25-a835-4cbd-b8a2-1418ba2c1d31",
   "metadata": {},
   "outputs": [],
   "source": [
    "# In the original dataset, `'AdoptionSpeed'` of `4` indicates\n",
    "# a pet was not adopted.\n",
    "dataframe['target'] = np.where(dataframe['AdoptionSpeed']==4, 0, 1)\n",
    "\n",
    "# Drop unused features.\n",
    "dataframe = dataframe.drop(columns=['AdoptionSpeed', 'Description'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00d403cf-9ae7-4780-9fac-13d920d8b395",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/rishic/anaconda3/envs/spark-dl-tf/lib/python3.11/site-packages/numpy/core/fromnumeric.py:59: FutureWarning: 'DataFrame.swapaxes' is deprecated and will be removed in a future version. Please use 'DataFrame.transpose' instead.\n",
      "  return bound(*args, **kwds)\n"
     ]
    }
   ],
   "source": [
    "train, val, test = np.split(dataframe.sample(frac=1), [int(0.8*len(dataframe)), int(0.9*len(dataframe))])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4206a56e-5403-42a9-805e-e037044e7995",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "9229 training examples\n",
      "1154 validation examples\n",
      "1154 test examples\n"
     ]
    }
   ],
   "source": [
    "print(len(train), 'training examples')\n",
    "print(len(val), 'validation examples')\n",
    "print(len(test), 'test examples')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7fa64f8",
   "metadata": {},
   "source": [
    "Create an input pipeline which converts each dataset into a tf.data.Dataset with shuffling and batching."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "499ade5f-ac8a-47ca-a021-071239dfe97d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def df_to_dataset(dataframe, shuffle=True, batch_size=32):\n",
    "    df = dataframe.copy()\n",
    "    labels = df.pop('target')\n",
    "    df = {key: value.to_numpy()[:,tf.newaxis] for key, value in dataframe.items()}\n",
    "    ds = tf.data.Dataset.from_tensor_slices((dict(df), labels))\n",
    "    if shuffle:\n",
    "        ds = ds.shuffle(buffer_size=len(dataframe))\n",
    "    ds = ds.batch(batch_size)\n",
    "    ds = ds.prefetch(batch_size)\n",
    "    return ds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96065bed",
   "metadata": {},
   "source": [
    "Check the format of the data returned by the pipeline:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9ec57c9-080e-4626-9e03-acf309cf3736",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0000 00:00:1738706370.981571 3686377 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "I0000 00:00:1738706370.984478 3686377 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "I0000 00:00:1738706370.987280 3686377 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "I0000 00:00:1738706371.105121 3686377 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "I0000 00:00:1738706371.106231 3686377 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "I0000 00:00:1738706371.107182 3686377 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2025-02-04 13:59:31.108098: I tensorflow/core/common_runtime/gpu/gpu_device.cc:2021] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 40337 MB memory:  -> device: 0, name: NVIDIA RTX A6000, pci bus id: 0000:01:00.0, compute capability: 8.6\n"
     ]
    }
   ],
   "source": [
    "batch_size = 5\n",
    "train_ds = df_to_dataset(train, batch_size=batch_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdc8571c",
   "metadata": {},
   "source": [
    "(Note that OUT_OF_RANGE errors are safe to ignore: https://github.com/tensorflow/tensorflow/issues/62963)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfcbf268-4508-4eb8-abe1-acf1dbb97bd5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Every feature: ['Type', 'Age', 'Breed1', 'Gender', 'Color1', 'Color2', 'MaturitySize', 'FurLength', 'Vaccinated', 'Sterilized', 'Health', 'Fee', 'PhotoAmt', 'target']\n",
      "A batch of ages: tf.Tensor(\n",
      "[[ 4]\n",
      " [60]\n",
      " [24]\n",
      " [ 1]\n",
      " [ 2]], shape=(5, 1), dtype=int64)\n",
      "A batch of targets: tf.Tensor([1 1 1 1 1], shape=(5,), dtype=int64)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-02-04 13:59:31.170523: I tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence\n"
     ]
    }
   ],
   "source": [
    "[(train_features, label_batch)] = train_ds.take(1)\n",
    "print('Every feature:', list(train_features.keys()))\n",
    "print('A batch of ages:', train_features['Age'])\n",
    "print('A batch of targets:', label_batch )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5a2d10c",
   "metadata": {},
   "source": [
    "### Apply Keras preprocessing layers\n",
    "\n",
    "We'll define a normalization layer for numeric features, and a category encoding for categorical features."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "6c09dc4b-3a2a-44f5-b41c-821ec30b87b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_normalization_layer(name, dataset):\n",
    "    # Create a Normalization layer for the feature.\n",
    "    normalizer = layers.Normalization(axis=None)\n",
    "\n",
    "    # Prepare a Dataset that only yields the feature.\n",
    "    feature_ds = dataset.map(lambda x, y: x[name])\n",
    "\n",
    "    # Learn the statistics of the data.\n",
    "    normalizer.adapt(feature_ds)\n",
    "\n",
    "    return normalizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "59bb91dc-360a-4a89-a9ea-bebc1ddbf1b7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-02-04 13:59:32.726183: I tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(5, 1), dtype=float32, numpy=\n",
       "array([[-0.19333968],\n",
       "       [-0.19333968],\n",
       "       [-0.19333968],\n",
       "       [-0.51676387],\n",
       "       [ 1.100357  ]], dtype=float32)>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "photo_count_col = train_features['PhotoAmt']\n",
    "layer = get_normalization_layer('PhotoAmt', train_ds)\n",
    "layer(photo_count_col)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "4623b612-e924-472b-9ef4-c7f14f9f53c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_category_encoding_layer(name, dataset, dtype, max_tokens=None):\n",
    "    # Create a layer that turns strings into integer indices.\n",
    "    if dtype == 'string':\n",
    "        index = layers.StringLookup(max_tokens=max_tokens)\n",
    "    # Otherwise, create a layer that turns integer values into integer indices.\n",
    "    else:\n",
    "        index = layers.IntegerLookup(max_tokens=max_tokens)\n",
    "\n",
    "    # Prepare a `tf.data.Dataset` that only yields the feature.\n",
    "    feature_ds = dataset.map(lambda x, y: x[name])\n",
    "\n",
    "    # Learn the set of possible values and assign them a fixed integer index.\n",
    "    index.adapt(feature_ds)\n",
    "\n",
    "    # Encode the integer indices.\n",
    "    encoder = layers.CategoryEncoding(num_tokens=index.vocabulary_size())\n",
    "\n",
    "    # Apply multi-hot encoding to the indices. The lambda function captures the\n",
    "    # layer, so you can use them, or include them in the Keras Functional model later.\n",
    "    return lambda feature: encoder(index(feature))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "0a40e9ee-20a5-4a42-8543-c267f99af55e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(5, 3), dtype=float32, numpy=\n",
       "array([[0., 0., 1.],\n",
       "       [0., 1., 0.],\n",
       "       [0., 0., 1.],\n",
       "       [0., 1., 0.],\n",
       "       [0., 1., 0.]], dtype=float32)>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_type_col = train_features['Type']\n",
    "test_type_layer = get_category_encoding_layer(name='Type',\n",
    "                                              dataset=train_ds,\n",
    "                                              dtype='string')\n",
    "test_type_layer(test_type_col)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "ff63a5cc-71f4-428e-9299-a8018edc7648",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-02-04 13:59:34.294276: I tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(5, 5), dtype=float32, numpy=\n",
       "array([[0., 0., 0., 0., 1.],\n",
       "       [1., 0., 0., 0., 0.],\n",
       "       [1., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 1., 0.],\n",
       "       [0., 1., 0., 0., 0.]], dtype=float32)>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_age_col = train_features['Age']\n",
    "test_age_layer = get_category_encoding_layer(name='Age',\n",
    "                                             dataset=train_ds,\n",
    "                                             dtype='int64',\n",
    "                                             max_tokens=5)\n",
    "test_age_layer(test_age_col)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afefbcf2",
   "metadata": {},
   "source": [
    "### Preprocess selected features\n",
    "\n",
    "Apply the preprocessing helper class defined earlier. Add all the feature inputs to a list.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "2b040b0e-d8ca-4cf0-917c-dd9a272e1f0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 256\n",
    "train_ds = df_to_dataset(train, batch_size=batch_size)\n",
    "val_ds = df_to_dataset(val, shuffle=False, batch_size=batch_size)\n",
    "test_ds = df_to_dataset(test, shuffle=False, batch_size=batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "19df498e-4dd1-467a-8741-e1f5e15932a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_inputs = {}\n",
    "encoded_features = []\n",
    "\n",
    "# Numerical features.\n",
    "for header in ['PhotoAmt', 'Fee']:\n",
    "    numeric_col = tf.keras.Input(shape=(1,), name=header)\n",
    "    normalization_layer = get_normalization_layer(header, train_ds)\n",
    "    encoded_numeric_col = normalization_layer(numeric_col)\n",
    "    all_inputs[header] = numeric_col\n",
    "    encoded_features.append(encoded_numeric_col)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "1d12579f-34fb-40b0-a16a-3e13cfea8178",
   "metadata": {},
   "outputs": [],
   "source": [
    "age_col = tf.keras.Input(shape=(1,), name='Age', dtype='int64')\n",
    "\n",
    "encoding_layer = get_category_encoding_layer(name='Age',\n",
    "                                             dataset=train_ds,\n",
    "                                             dtype='int64',\n",
    "                                             max_tokens=5)\n",
    "encoded_age_col = encoding_layer(age_col)\n",
    "all_inputs['Age'] = age_col\n",
    "encoded_features.append(encoded_age_col)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "bff286eb-7ad7-4d3a-8fa4-c729692d1425",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-02-04 13:59:34.588989: I tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence\n",
      "2025-02-04 13:59:35.029267: I tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence\n"
     ]
    }
   ],
   "source": [
    "categorical_cols = ['Type', 'Color1', 'Color2', 'Gender', 'MaturitySize',\n",
    "                    'FurLength', 'Vaccinated', 'Sterilized', 'Health', 'Breed1']\n",
    "\n",
    "for header in categorical_cols:\n",
    "    categorical_col = tf.keras.Input(shape=(1,), name=header, dtype='string')\n",
    "    encoding_layer = get_category_encoding_layer(name=header,\n",
    "                                                dataset=train_ds,\n",
    "                                                dtype='string',\n",
    "                                                max_tokens=5)\n",
    "    encoded_categorical_col = encoding_layer(categorical_col)\n",
    "    all_inputs[header] = categorical_col\n",
    "    encoded_features.append(encoded_categorical_col)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0dfac0d",
   "metadata": {},
   "source": [
    "### Create, compile, and train model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "79247436-32d8-4738-a656-3f288c77001c",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_features = tf.keras.layers.concatenate(encoded_features)\n",
    "x = tf.keras.layers.Dense(32, activation=\"relu\")(all_features)\n",
    "x = tf.keras.layers.Dropout(0.5)(x)\n",
    "output = tf.keras.layers.Dense(1)(x)\n",
    "\n",
    "model = tf.keras.Model(all_inputs, output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "dbc85d3e-6d1e-4167-9516-b1182e880542",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.compile(optimizer='adam',\n",
    "              loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),\n",
    "              metrics=[\"accuracy\"],\n",
    "              run_eagerly=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "bc9836c8-3c1a-41ad-8833-a946bafcfb00",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/10\n",
      "\u001b[1m37/37\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m1s\u001b[0m 16ms/step - accuracy: 0.3658 - loss: 0.7746 - val_accuracy: 0.6854 - val_loss: 0.5841\n",
      "Epoch 2/10\n",
      "\u001b[1m37/37\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m1s\u001b[0m 16ms/step - accuracy: 0.6270 - loss: 0.6023 - val_accuracy: 0.7383 - val_loss: 0.5593\n",
      "Epoch 3/10\n",
      "\u001b[1m37/37\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m1s\u001b[0m 16ms/step - accuracy: 0.6650 - loss: 0.5781 - val_accuracy: 0.7392 - val_loss: 0.5442\n",
      "Epoch 4/10\n",
      "\u001b[1m37/37\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m1s\u001b[0m 17ms/step - accuracy: 0.6609 - loss: 0.5744 - val_accuracy: 0.7418 - val_loss: 0.5329\n",
      "Epoch 5/10\n",
      "\u001b[1m37/37\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m1s\u001b[0m 15ms/step - accuracy: 0.6845 - loss: 0.5555 - val_accuracy: 0.7444 - val_loss: 0.5261\n",
      "Epoch 6/10\n",
      "\u001b[1m37/37\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m1s\u001b[0m 15ms/step - accuracy: 0.6910 - loss: 0.5465 - val_accuracy: 0.7513 - val_loss: 0.5198\n",
      "Epoch 7/10\n",
      "\u001b[1m37/37\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m1s\u001b[0m 15ms/step - accuracy: 0.7018 - loss: 0.5475 - val_accuracy: 0.7556 - val_loss: 0.5145\n",
      "Epoch 8/10\n",
      "\u001b[1m37/37\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m1s\u001b[0m 15ms/step - accuracy: 0.7026 - loss: 0.5410 - val_accuracy: 0.7496 - val_loss: 0.5099\n",
      "Epoch 9/10\n",
      "\u001b[1m37/37\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m1s\u001b[0m 15ms/step - accuracy: 0.7145 - loss: 0.5315 - val_accuracy: 0.7530 - val_loss: 0.5066\n",
      "Epoch 10/10\n",
      "\u001b[1m37/37\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m1s\u001b[0m 15ms/step - accuracy: 0.7099 - loss: 0.5316 - val_accuracy: 0.7539 - val_loss: 0.5038\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<keras.src.callbacks.history.History at 0x7b15e89123d0>"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.fit(train_ds, epochs=10, validation_data=val_ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbccebaa-fc24-4a58-a032-222cef8fdf08",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m5/5\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m0s\u001b[0m 7ms/step - accuracy: 0.7416 - loss: 0.5196 \n",
      "Accuracy 0.7443674206733704\n"
     ]
    }
   ],
   "source": [
    "loss, accuracy = model.evaluate(test_ds)\n",
    "print(\"Accuracy\", accuracy)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7534616c-8561-4869-b6e9-7254ebdb2c3f",
   "metadata": {},
   "source": [
    "### Save and reload model\n",
    "\n",
    "Demonstrate saving the trained model and reloading it for inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "6bf0d024",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save('models/my_pet_classifier.keras')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "d1a7be62",
   "metadata": {},
   "outputs": [],
   "source": [
    "reloaded_model = tf.keras.models.load_model('models/my_pet_classifier.keras')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3d2a2d5-fd4d-4320-bacc-fd4571cec709",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m1/1\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m0s\u001b[0m 27ms/step\n",
      "This particular pet had a 83.2 percent probability of getting adopted.\n"
     ]
    }
   ],
   "source": [
    "sample = {\n",
    "    'Type': 'Cat',\n",
    "    'Age': 3,\n",
    "    'Breed1': 'Tabby',\n",
    "    'Gender': 'Male',\n",
    "    'Color1': 'Black',\n",
    "    'Color2': 'White',\n",
    "    'MaturitySize': 'Small',\n",
    "    'FurLength': 'Short',\n",
    "    'Vaccinated': 'No',\n",
    "    'Sterilized': 'No',\n",
    "    'Health': 'Healthy',\n",
    "    'Fee': 100,\n",
    "    'PhotoAmt': 2,\n",
    "}\n",
    "\n",
    "input_dict = {name: tf.convert_to_tensor([value]) for name, value in sample.items()}\n",
    "predictions = reloaded_model.predict(input_dict)\n",
    "prob = tf.nn.sigmoid(predictions[0])\n",
    "\n",
    "print(\n",
    "    \"This particular pet had a %.1f percent probability \"\n",
    "    \"of getting adopted.\" % (100 * prob)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7bbfe69-93ed-4452-8985-c6685e0726c3",
   "metadata": {},
   "source": [
    "## PySpark"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "fc8a0536",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, struct, pandas_udf\n",
    "from pyspark.ml.functions import predict_batch_udf\n",
    "from pyspark.sql.types import *\n",
    "from pyspark.sql import SparkSession\n",
    "from pyspark import SparkConf\n",
    "import json\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb5aa875",
   "metadata": {},
   "source": [
    "Check the cluster environment to handle any platform-specific Spark configurations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "7701420e",
   "metadata": {},
   "outputs": [],
   "source": [
    "on_databricks = os.environ.get(\"DATABRICKS_RUNTIME_VERSION\", False)\n",
    "on_dataproc = os.environ.get(\"DATAPROC_IMAGE_VERSION\", False)\n",
    "on_standalone = not (on_databricks or on_dataproc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e231dbd",
   "metadata": {},
   "source": [
    "#### Create Spark Session\n",
    "\n",
    "For local standalone clusters, we'll connect to the cluster and create the Spark Session.  \n",
    "For CSP environments, Spark will either be preconfigured (Databricks) or we'll need to create the Spark Session (Dataproc)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "60dff1da",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "25/02/04 13:59:42 WARN Utils: Your hostname, cb4ae00-lcedt resolves to a loopback address: 127.0.1.1; using 10.110.47.100 instead (on interface eno1)\n",
      "25/02/04 13:59:42 WARN Utils: Set SPARK_LOCAL_IP if you need to bind to another address\n",
      "Setting default log level to \"WARN\".\n",
      "To adjust logging level use sc.setLogLevel(newLevel). For SparkR, use setLogLevel(newLevel).\n",
      "25/02/04 13:59:43 WARN NativeCodeLoader: Unable to load native-hadoop library for your platform... using builtin-java classes where applicable\n"
     ]
    }
   ],
   "source": [
    "conf = SparkConf()\n",
    "\n",
    "if 'spark' not in globals():\n",
    "    if on_standalone:\n",
    "        import socket\n",
    "        \n",
    "        conda_env = os.environ.get(\"CONDA_PREFIX\")\n",
    "        hostname = socket.gethostname()\n",
    "        conf.setMaster(f\"spark://{hostname}:7077\")\n",
    "        conf.set(\"spark.pyspark.python\", f\"{conda_env}/bin/python\")\n",
    "        conf.set(\"spark.pyspark.driver.python\", f\"{conda_env}/bin/python\")\n",
    "    elif on_dataproc:\n",
    "        conf.set(\"spark.executorEnv.TF_GPU_ALLOCATOR\", \"cuda_malloc_async\")\n",
    "\n",
    "    conf.set(\"spark.executor.cores\", \"8\")\n",
    "    conf.set(\"spark.task.resource.gpu.amount\", \"0.125\")\n",
    "    conf.set(\"spark.executor.resource.gpu.amount\", \"1\")\n",
    "    conf.set(\"spark.sql.execution.arrow.pyspark.enabled\", \"true\")\n",
    "    conf.set(\"spark.python.worker.reuse\", \"true\")\n",
    "\n",
    "conf.set(\"spark.sql.execution.arrow.maxRecordsPerBatch\", \"1000\")\n",
    "spark = SparkSession.builder.appName(\"spark-dl-examples\").config(conf=conf).getOrCreate()\n",
    "sc = spark.sparkContext"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa2333d1",
   "metadata": {},
   "source": [
    "#### Create PySpark DataFrame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "3c64fd7b-3d1e-40f8-ab64-b5c13f8bbe77",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = spark.createDataFrame(dataframe).repartition(8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "1be8215b-5068-41b4-849c-1c3ea7bb108a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "data_path = \"spark-dl-datasets/petfinder-mini\"\n",
    "if on_databricks:\n",
    "    dbutils.fs.mkdirs(\"/FileStore/spark-dl-datasets\")\n",
    "    data_path = \"dbfs:/FileStore/\" + data_path\n",
    "\n",
    "df.write.mode(\"overwrite\").parquet(data_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cec4e0e",
   "metadata": {},
   "source": [
    "#### Load and preprocess DataFrame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "0892f845",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = spark.read.parquet(data_path).cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "952645dd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Type', 'Age', 'Breed1', 'Gender', 'Color1', 'Color2', 'MaturitySize', 'FurLength', 'Vaccinated', 'Sterilized', 'Health', 'Fee', 'PhotoAmt', 'target']\n"
     ]
    }
   ],
   "source": [
    "columns = df.columns\n",
    "print(columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "b9c24c0d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Type', 'Age', 'Breed1', 'Gender', 'Color1', 'Color2', 'MaturitySize', 'FurLength', 'Vaccinated', 'Sterilized', 'Health', 'Fee', 'PhotoAmt']\n"
     ]
    }
   ],
   "source": [
    "# remove label column\n",
    "columns.remove(\"target\")\n",
    "print(columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "d4dbde99-cf65-4c15-a163-754a0201a48d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----+---+--------------------+------+------+--------+------------+---------+----------+----------+-------+---+--------+------+\n",
      "|Type|Age|              Breed1|Gender|Color1|  Color2|MaturitySize|FurLength|Vaccinated|Sterilized| Health|Fee|PhotoAmt|target|\n",
      "+----+---+--------------------+------+------+--------+------------+---------+----------+----------+-------+---+--------+------+\n",
      "| Dog|  3|         Mixed Breed|  Male| Black|No Color|       Small|   Medium|  Not Sure|  Not Sure|Healthy|  0|       2|     0|\n",
      "| Dog|  9|         Mixed Breed|  Male|  Gray|No Color|      Medium|    Short|  Not Sure|        No|Healthy|  0|       4|     1|\n",
      "| Cat|  4| Domestic Short Hair|  Male| Black|    Gray|      Medium|    Short|  Not Sure|  Not Sure|Healthy|  0|       4|     1|\n",
      "| Cat|  6| Domestic Short Hair|  Male|Yellow|   White|      Medium|    Short|        No|        No|Healthy|  0|       3|     1|\n",
      "| Cat|  6|Domestic Medium Hair|  Male|  Gray|No Color|       Small|   Medium|        No|        No|Healthy|  0|       4|     1|\n",
      "+----+---+--------------------+------+------+--------+------------+---------+----------+----------+-------+---+--------+------+\n",
      "only showing top 5 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df.show(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "824d7f97",
   "metadata": {},
   "source": [
    "## Inference using Spark DL API\n",
    "\n",
    "Distributed inference using the PySpark [predict_batch_udf](https://spark.apache.org/docs/latest/api/python/reference/api/pyspark.ml.functions.predict_batch_udf.html#pyspark.ml.functions.predict_batch_udf):\n",
    "\n",
    "- predict_batch_fn uses Tensorflow APIs to load the model and return a predict function which operates on numpy arrays \n",
    "- predict_batch_udf will convert the Spark DataFrame columns into numpy input batches for the predict function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "d62eb95a-54c6-44d2-9279-38fb65e0e160",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get absolute path to model\n",
    "model_path = \"{}/models/my_pet_classifier.keras\".format(os.getcwd())\n",
    "\n",
    "# For cloud environments, copy the model to the distributed file system.\n",
    "if on_databricks:\n",
    "    dbutils.fs.mkdirs(\"/FileStore/spark-dl-models\")\n",
    "    dbfs_model_path = \"/dbfs/FileStore/spark-dl-models/my_pet_classifier.keras\"\n",
    "    shutil.copy(model_path, dbfs_model_path)\n",
    "    model_path = dbfs_model_path\n",
    "elif on_dataproc:\n",
    "    # GCS is mounted at /mnt/gcs by the init script\n",
    "    models_dir = \"/mnt/gcs/spark-dl/models\"\n",
    "    os.mkdir(models_dir) if not os.path.exists(models_dir) else None\n",
    "    gcs_model_path = models_dir + \"/my_pet_classifier.keras\"\n",
    "    shutil.copy(model_path, gcs_model_path)\n",
    "    model_path = gcs_model_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "45665acf-50c8-445b-a985-b3dabd734709",
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_batch_fn():\n",
    "    import tensorflow as tf\n",
    "    import pandas as pd\n",
    "    \n",
    "    # Enable GPU memory growth to avoid CUDA OOM\n",
    "    gpus = tf.config.experimental.list_physical_devices('GPU')\n",
    "    if gpus:\n",
    "        try:\n",
    "            for gpu in gpus:\n",
    "                tf.config.experimental.set_memory_growth(gpu, True)\n",
    "        except RuntimeError as e:\n",
    "            print(e)\n",
    "\n",
    "    model = tf.keras.models.load_model(model_path)\n",
    "\n",
    "    def predict(t, a, b, g, c1, c2, m, f, v, s, h, fee, p):\n",
    "        inputs = {\n",
    "            \"Type\": t,\n",
    "            \"Age\": a,\n",
    "            \"Breed1\": b,\n",
    "            \"Gender\": g,\n",
    "            \"Color1\": c1,\n",
    "            \"Color2\": c2,\n",
    "            \"MaturitySize\": m,\n",
    "            \"FurLength\": f,\n",
    "            \"Vaccinated\": v,\n",
    "            \"Sterilized\": s,\n",
    "            \"Health\": h,\n",
    "            \"Fee\": fee,\n",
    "            \"PhotoAmt\": p\n",
    "        }\n",
    "        # return model.predict(inputs)\n",
    "        return pd.Series(np.squeeze(model.predict(inputs)))\n",
    "\n",
    "    return predict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "815e3b5f-7914-4235-85fa-50153dcd3d30",
   "metadata": {},
   "outputs": [],
   "source": [
    "# need to pass the list of columns into the model_udf\n",
    "classify = predict_batch_udf(predict_batch_fn,\n",
    "                             return_type=FloatType(),\n",
    "                             batch_size=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "da03a0c6-2d39-425e-a9fa-57c139cca1ed",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "25/02/04 13:59:47 WARN package: Truncated the string representation of a plan since it was too large. This behavior can be adjusted by setting 'spark.sql.debug.maxToStringFields'.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 5:>                                                          (0 + 8) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 19.8 ms, sys: 9.3 ms, total: 29.1 ms\n",
      "Wall time: 4.99 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "preds = df.withColumn(\"preds\", classify(struct(*columns)))\n",
    "results = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "03990c76-7198-49a7-bb5d-6870be915fb3",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 6:>                                                          (0 + 8) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 86.9 ms, sys: 13.7 ms, total: 101 ms\n",
      "Wall time: 1.56 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "preds = df.withColumn(\"preds\", classify(*columns))\n",
    "results = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "edb93cf3-c248-40c9-b8dc-acc8f51786a9",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 7:>                                                          (0 + 8) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 16.4 ms, sys: 4.46 ms, total: 20.9 ms\n",
      "Wall time: 1.52 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "preds = df.withColumn(\"preds\", classify(*[col(c) for c in columns]))\n",
    "results = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "a91f19cb-f7f1-4669-aff1-be594bea5378",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----+---+--------------------+------+------+--------+------------+---------+----------+----------+------------+---+--------+------+-----------+\n",
      "|Type|Age|              Breed1|Gender|Color1|  Color2|MaturitySize|FurLength|Vaccinated|Sterilized|      Health|Fee|PhotoAmt|target|      preds|\n",
      "+----+---+--------------------+------+------+--------+------------+---------+----------+----------+------------+---+--------+------+-----------+\n",
      "| Dog|  3|         Mixed Breed|  Male| Black|No Color|       Small|   Medium|  Not Sure|  Not Sure|     Healthy|  0|       2|     0|  0.4963937|\n",
      "| Dog|  9|         Mixed Breed|  Male|  Gray|No Color|      Medium|    Short|  Not Sure|        No|     Healthy|  0|       4|     1|  0.6780287|\n",
      "| Cat|  4| Domestic Short Hair|  Male| Black|    Gray|      Medium|    Short|  Not Sure|  Not Sure|     Healthy|  0|       4|     1| 0.58800673|\n",
      "| Cat|  6| Domestic Short Hair|  Male|Yellow|   White|      Medium|    Short|        No|        No|     Healthy|  0|       3|     1|  0.7378843|\n",
      "| Cat|  6|Domestic Medium Hair|  Male|  Gray|No Color|       Small|   Medium|        No|        No|     Healthy|  0|       4|     1|  1.2695599|\n",
      "| Cat|  5|Domestic Medium Hair|Female|  Gray|No Color|      Medium|   Medium|       Yes|  Not Sure|     Healthy|  0|       1|     0|0.060457088|\n",
      "| Dog| 24|              Beagle|Female| Black|  Golden|      Medium|    Short|  Not Sure|  Not Sure|Minor Injury|  0|       1|     1| 0.28160828|\n",
      "| Cat| 29|               Tabby|  Male| Brown|  Golden|      Medium|    Short|        No|        No|     Healthy|  0|       1|     0|  0.6928505|\n",
      "| Dog|  9|         Mixed Breed|Female| Black|   Brown|      Medium|    Short|       Yes|       Yes|     Healthy|  0|       2|     0|-0.10125986|\n",
      "| Dog|  2|         Mixed Breed|Female| Cream|   White|      Medium|    Short|        No|        No|     Healthy|  0|       1|     0|  1.3703903|\n",
      "| Dog|  2|         Mixed Breed|  Male| Brown|   White|      Medium|    Short|       Yes|        No|     Healthy|  0|       1|     1|  1.3243997|\n",
      "| Dog| 60|    Golden Retriever|  Male| Brown|  Yellow|      Medium|   Medium|       Yes|       Yes|     Healthy|  0|       5|     1|  0.9026731|\n",
      "| Cat|  9|             Siamese|  Male| White|No Color|      Medium|    Short|       Yes|        No|     Healthy|  0|       2|     1|  0.8207382|\n",
      "| Dog| 19|   Doberman Pinscher|Female| Black|   Brown|       Large|    Short|       Yes|       Yes|     Healthy|500|       2|     1| 0.85343015|\n",
      "| Cat| 11| Domestic Short Hair|  Male| Cream|No Color|      Medium|    Short|       Yes|       Yes|     Healthy|100|       6|     0| 0.53920615|\n",
      "| Dog| 18|         Mixed Breed|Female| Brown|   White|       Small|    Short|       Yes|        No|     Healthy|  0|       5|     0|   0.718272|\n",
      "| Dog|  4|         Mixed Breed|Female| Brown|   White|      Medium|   Medium|  Not Sure|  Not Sure|     Healthy|  0|       3|     0| 0.16185221|\n",
      "| Dog| 96|    Golden Retriever|  Male|Golden|No Color|       Large|     Long|       Yes|       Yes|     Healthy|  0|       2|     1|  0.8156965|\n",
      "| Dog| 54|    Golden Retriever|  Male|Golden|No Color|       Large|   Medium|       Yes|        No|     Healthy|350|      20|     1|  3.5315154|\n",
      "| Cat|  5|Domestic Medium Hair|Female| Brown|   White|      Medium|   Medium|        No|        No|     Healthy|  0|       5|     1|  1.1725564|\n",
      "+----+---+--------------------+------+------+--------+------------+---------+----------+----------+------------+---+--------+------+-----------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "preds.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c3e0390",
   "metadata": {},
   "source": [
    "## Using Triton Inference Server\n",
    "In this section, we demonstrate integration with the [Triton Inference Server](https://developer.nvidia.com/nvidia-triton-inference-server), an open-source, GPU-accelerated serving solution for DL.  \n",
    "We use [PyTriton](https://github.com/triton-inference-server/pytriton), a Flask-like framework that handles client/server communication with the Triton server.  \n",
    "\n",
    "The process looks like this:\n",
    "- Distribute a PyTriton task across the Spark cluster, instructing each node to launch a Triton server process.\n",
    "- Define a Triton inference function, which contains a client that binds to the local server on a given node and sends inference requests.\n",
    "- Wrap the Triton inference function in a predict_batch_udf to launch parallel inference requests using Spark.\n",
    "- Finally, distribute a shutdown signal to terminate the Triton server processes on each node.\n",
    "\n",
    "<img src=\"../images/spark-server.png\" alt=\"drawing\" width=\"700\"/>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "2605d134-ef75-4d94-9b16-2c6d85f29bef",
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea407357",
   "metadata": {},
   "source": [
    "Import the helper class from server_utils.py:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "7e1e716f",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.addPyFile(\"server_utils.py\")\n",
    "\n",
    "from server_utils import TritonServerManager"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcd28e7d",
   "metadata": {},
   "source": [
    "Define the Triton Server function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "4666e618-8038-4dc5-9be7-793aedbf4500",
   "metadata": {},
   "outputs": [],
   "source": [
    "def triton_server(ports, model_path):\n",
    "    import time\n",
    "    import signal\n",
    "    import numpy as np\n",
    "    import tensorflow as tf\n",
    "    from pytriton.decorators import batch\n",
    "    from pytriton.model_config import DynamicBatcher, ModelConfig, Tensor\n",
    "    from pytriton.triton import Triton, TritonConfig\n",
    "    from pyspark import TaskContext\n",
    "\n",
    "    print(f\"SERVER: Initializing model on worker {TaskContext.get().partitionId()}.\")\n",
    "    # Enable GPU memory growth\n",
    "    gpus = tf.config.experimental.list_physical_devices('GPU')\n",
    "    if gpus:\n",
    "        try:\n",
    "            for gpu in gpus:\n",
    "                tf.config.experimental.set_memory_growth(gpu, True)\n",
    "        except RuntimeError as e:\n",
    "            print(e)\n",
    "\n",
    "    model = tf.keras.models.load_model(model_path)\n",
    "\n",
    "    def decode(input_tensor):\n",
    "        return tf.convert_to_tensor(np.vectorize(lambda x: x.decode('utf-8'))(input_tensor))\n",
    "\n",
    "    def identity(input_tensor):\n",
    "        return tf.convert_to_tensor(input_tensor)\n",
    "\n",
    "    input_transforms = {\n",
    "        \"Type\": decode,\n",
    "        \"Age\": identity,\n",
    "        \"Breed1\": decode,\n",
    "        \"Gender\": decode,\n",
    "        \"Color1\": decode,\n",
    "        \"Color2\": decode,\n",
    "        \"MaturitySize\": decode,\n",
    "        \"FurLength\": decode,\n",
    "        \"Vaccinated\": decode,\n",
    "        \"Sterilized\": decode,\n",
    "        \"Health\": decode,\n",
    "        \"Fee\": identity,\n",
    "        \"PhotoAmt\": identity\n",
    "    }\n",
    "\n",
    "    @batch\n",
    "    def _infer_fn(**inputs):\n",
    "        decoded_inputs = {k: input_transforms[k](v) for k, v in inputs.items()}\n",
    "        print(f\"SERVER: Received batch of size {len(decoded_inputs['Type'])}.\")\n",
    "        return {\n",
    "            \"preds\": model.predict(decoded_inputs)\n",
    "        }\n",
    "\n",
    "    workspace_path = f\"/tmp/triton_{time.strftime('%m_%d_%M_%S')}\"\n",
    "    triton_conf = TritonConfig(http_port=ports[0], grpc_port=ports[1], metrics_port=ports[2])\n",
    "    with Triton(config=triton_conf, workspace=workspace_path) as triton:\n",
    "        triton.bind(\n",
    "            model_name=\"PetClassifier\",\n",
    "            infer_func=_infer_fn,\n",
    "            inputs=[\n",
    "                Tensor(name=\"Type\", dtype=np.bytes_, shape=(-1,)),\n",
    "                Tensor(name=\"Age\", dtype=np.int64, shape=(-1,)),\n",
    "                Tensor(name=\"Breed1\", dtype=np.bytes_, shape=(-1,)),\n",
    "                Tensor(name=\"Gender\", dtype=np.bytes_, shape=(-1,)),\n",
    "                Tensor(name=\"Color1\", dtype=np.bytes_, shape=(-1,)),\n",
    "                Tensor(name=\"Color2\", dtype=np.bytes_, shape=(-1,)),\n",
    "                Tensor(name=\"MaturitySize\", dtype=np.bytes_, shape=(-1,)),\n",
    "                Tensor(name=\"FurLength\", dtype=np.bytes_, shape=(-1,)),\n",
    "                Tensor(name=\"Vaccinated\", dtype=np.bytes_, shape=(-1,)),\n",
    "                Tensor(name=\"Sterilized\", dtype=np.bytes_, shape=(-1,)),\n",
    "                Tensor(name=\"Health\", dtype=np.bytes_, shape=(-1,)),\n",
    "                Tensor(name=\"Fee\", dtype=np.int64, shape=(-1,)),\n",
    "                Tensor(name=\"PhotoAmt\", dtype=np.int64, shape=(-1,)),\n",
    "            ],\n",
    "            outputs=[\n",
    "                Tensor(name=\"preds\", dtype=np.float32, shape=(-1,)),\n",
    "            ],\n",
    "            config=ModelConfig(\n",
    "                max_batch_size=128,\n",
    "                batcher=DynamicBatcher(max_queue_delay_microseconds=5000),  # 5ms\n",
    "            ),\n",
    "            strict=True,\n",
    "        )\n",
    "\n",
    "        def _stop_triton(signum, frame):\n",
    "            # The server manager sends SIGTERM to stop the server; this function ensures graceful cleanup.\n",
    "            print(\"SERVER: Received SIGTERM. Stopping Triton server.\")\n",
    "            triton.stop()\n",
    "\n",
    "        signal.signal(signal.SIGTERM, _stop_triton)\n",
    "\n",
    "        print(\"SERVER: Serving inference\")\n",
    "        triton.serve()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "617525a5",
   "metadata": {},
   "source": [
    "#### Start Triton servers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc93a43a",
   "metadata": {},
   "source": [
    "The `TritonServerManager` will handle the lifecycle of Triton server instances across the Spark cluster:\n",
    "- Find available ports for HTTP/gRPC/metrics\n",
    "- Deploy a server on each node via stage-level scheduling\n",
    "- Gracefully shutdown servers across nodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9b98208",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"PetClassifier\"\n",
    "server_manager = TritonServerManager(model_name=model_name, model_path=model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "228401f7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-02-07 11:03:44,809 - INFO - Requesting stage-level resources: (cores=5, gpu=1.0)\n",
      "2025-02-07 11:03:44,810 - INFO - Starting 1 servers.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'cb4ae00-lcedt': (2020631, [7000, 7001, 7002])}"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Returns {'hostname', (server_pid, [http_port, grpc_port, metrics_port])}\n",
    "server_manager.start_servers(triton_server)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb560288",
   "metadata": {},
   "source": [
    "#### Define client function"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d28b1ca",
   "metadata": {},
   "source": [
    "Get the hostname -> url mapping from the server manager:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1234a02",
   "metadata": {},
   "outputs": [],
   "source": [
    "host_to_http_url = server_manager.host_to_http_url  # or server_manager.host_to_grpc_url"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c9ef706",
   "metadata": {},
   "source": [
    "Define the Triton inference function, which returns a predict function for batch inference through the server:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "e50b5fc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def triton_fn(model_name, host_to_url):\n",
    "    import socket\n",
    "    import numpy as np\n",
    "    from pytriton.client import ModelClient\n",
    "\n",
    "    url = host_to_url[socket.gethostname()]\n",
    "    print(f\"CLIENT: Connecting to {model_name} at {url}\")\n",
    "\n",
    "    def infer_batch(t, a, b, g, c1, c2, m, f, v, s, h, fee, p):\n",
    "        \n",
    "        def encode(value):\n",
    "            return np.vectorize(lambda x: x.encode(\"utf-8\"))(value).astype(np.bytes_)\n",
    "\n",
    "        with ModelClient(url, model_name, inference_timeout_s=240) as client:\n",
    "            encoded_inputs = {\n",
    "                \"Type\": encode(t), \n",
    "                \"Age\": a, \n",
    "                \"Breed1\": encode(b), \n",
    "                \"Gender\": encode(g),\n",
    "                \"Color1\": encode(c1),\n",
    "                \"Color2\": encode(c2),\n",
    "                \"MaturitySize\": encode(m),\n",
    "                \"FurLength\": encode(f),\n",
    "                \"Vaccinated\": encode(v),\n",
    "                \"Sterilized\": encode(s),\n",
    "                \"Health\": encode(h),\n",
    "                \"Fee\": fee,\n",
    "                \"PhotoAmt\": p\n",
    "            }\n",
    "            result_data = client.infer_batch(**encoded_inputs)\n",
    "            return result_data[\"preds\"]\n",
    "            \n",
    "    return infer_batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "2ffb020e-dc93-456b-bee6-405611eee1e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# need to pass the list of columns into the model_udf\n",
    "classify = predict_batch_udf(partial(triton_fn, model_name=model_name, host_to_url=host_to_http_url),\n",
    "                             input_tensor_shapes=[[1]] * len(columns),\n",
    "                             return_type=FloatType(),\n",
    "                             batch_size=64)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2edd887f",
   "metadata": {},
   "source": [
    "#### Load and preprocess DataFrame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "fe8dc3e6-f1b1-4a24-85f4-0a5ecabef4c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = spark.read.parquet(data_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "4cfb3f34-a215-4781-91bf-2bec85e15633",
   "metadata": {},
   "outputs": [],
   "source": [
    "columns = df.columns\n",
    "# remove label column\n",
    "columns.remove(\"target\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b75e6f20-f06c-4f4c-ada1-c562e078ed4b",
   "metadata": {},
   "source": [
    "#### Run inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "e6ff0356-becd-421f-aebb-272497d5ad6a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 12:>                                                         (0 + 8) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 17.3 ms, sys: 7.75 ms, total: 25.1 ms\n",
      "Wall time: 6.35 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "preds = df.withColumn(\"preds\", classify(struct(*columns)))\n",
    "results = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "ce18ee7c-5958-4986-b200-6d986fcc6243",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 15.2 ms, sys: 4.2 ms, total: 19.4 ms\n",
      "Wall time: 5.86 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "preds = df.withColumn(\"preds\", classify(*columns))\n",
    "results = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "0888ce40-b2c4-4aed-8ccb-6a8bcd00abc8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 14:===========================================>              (6 + 2) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 93.4 ms, sys: 3.4 ms, total: 96.8 ms\n",
      "Wall time: 5.87 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "preds = df.withColumn(\"preds\", classify(*[col(c) for c in columns]))\n",
    "results = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "d45812b5-f584-41a4-a821-2b59e065671c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----+---+--------------------+------+------+--------+------------+---------+----------+----------+------------+---+--------+------+-----------+\n",
      "|Type|Age|              Breed1|Gender|Color1|  Color2|MaturitySize|FurLength|Vaccinated|Sterilized|      Health|Fee|PhotoAmt|target|      preds|\n",
      "+----+---+--------------------+------+------+--------+------------+---------+----------+----------+------------+---+--------+------+-----------+\n",
      "| Dog|  3|         Mixed Breed|  Male| Black|No Color|       Small|   Medium|  Not Sure|  Not Sure|     Healthy|  0|       2|     0|  0.4963937|\n",
      "| Dog|  9|         Mixed Breed|  Male|  Gray|No Color|      Medium|    Short|  Not Sure|        No|     Healthy|  0|       4|     1|  0.6780287|\n",
      "| Cat|  4| Domestic Short Hair|  Male| Black|    Gray|      Medium|    Short|  Not Sure|  Not Sure|     Healthy|  0|       4|     1| 0.58800673|\n",
      "| Cat|  6| Domestic Short Hair|  Male|Yellow|   White|      Medium|    Short|        No|        No|     Healthy|  0|       3|     1|  0.7378843|\n",
      "| Cat|  6|Domestic Medium Hair|  Male|  Gray|No Color|       Small|   Medium|        No|        No|     Healthy|  0|       4|     1|  1.2695599|\n",
      "| Cat|  5|Domestic Medium Hair|Female|  Gray|No Color|      Medium|   Medium|       Yes|  Not Sure|     Healthy|  0|       1|     0|0.060457088|\n",
      "| Dog| 24|              Beagle|Female| Black|  Golden|      Medium|    Short|  Not Sure|  Not Sure|Minor Injury|  0|       1|     1| 0.28160828|\n",
      "| Cat| 29|               Tabby|  Male| Brown|  Golden|      Medium|    Short|        No|        No|     Healthy|  0|       1|     0|  0.6928505|\n",
      "| Dog|  9|         Mixed Breed|Female| Black|   Brown|      Medium|    Short|       Yes|       Yes|     Healthy|  0|       2|     0|-0.10125986|\n",
      "| Dog|  2|         Mixed Breed|Female| Cream|   White|      Medium|    Short|        No|        No|     Healthy|  0|       1|     0|  1.3703903|\n",
      "| Dog|  2|         Mixed Breed|  Male| Brown|   White|      Medium|    Short|       Yes|        No|     Healthy|  0|       1|     1|  1.3243997|\n",
      "| Dog| 60|    Golden Retriever|  Male| Brown|  Yellow|      Medium|   Medium|       Yes|       Yes|     Healthy|  0|       5|     1|  0.9026731|\n",
      "| Cat|  9|             Siamese|  Male| White|No Color|      Medium|    Short|       Yes|        No|     Healthy|  0|       2|     1|  0.8207382|\n",
      "| Dog| 19|   Doberman Pinscher|Female| Black|   Brown|       Large|    Short|       Yes|       Yes|     Healthy|500|       2|     1| 0.85343015|\n",
      "| Cat| 11| Domestic Short Hair|  Male| Cream|No Color|      Medium|    Short|       Yes|       Yes|     Healthy|100|       6|     0| 0.53920615|\n",
      "| Dog| 18|         Mixed Breed|Female| Brown|   White|       Small|    Short|       Yes|        No|     Healthy|  0|       5|     0|   0.718272|\n",
      "| Dog|  4|         Mixed Breed|Female| Brown|   White|      Medium|   Medium|  Not Sure|  Not Sure|     Healthy|  0|       3|     0| 0.16185221|\n",
      "| Dog| 96|    Golden Retriever|  Male|Golden|No Color|       Large|     Long|       Yes|       Yes|     Healthy|  0|       2|     1|  0.8156965|\n",
      "| Dog| 54|    Golden Retriever|  Male|Golden|No Color|       Large|   Medium|       Yes|        No|     Healthy|350|      20|     1|  3.5315154|\n",
      "| Cat|  5|Domestic Medium Hair|Female| Brown|   White|      Medium|   Medium|        No|        No|     Healthy|  0|       5|     1|  1.1725564|\n",
      "+----+---+--------------------+------+------+--------+------------+---------+----------+----------+------------+---+--------+------+-----------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "preds.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63135aa0-b44c-4dda-8050-8cad320afe88",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Stop Triton Server on each executor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "6914f44f-677f-4db3-be09-783df8d11b8a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-02-04 14:00:18,330 - INFO - Requesting stage-level resources: (cores=5, gpu=1.0)\n",
      "2025-02-04 14:00:28,520 - INFO - Sucessfully stopped 1 servers.                 \n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[True]"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "server_manager.stop_servers()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "f8c6ee43-8891-4446-986e-1447c5d48bac",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not on_databricks: # on databricks, spark.stop() puts the cluster in a bad state\n",
    "    spark.stop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e611126e-d8c3-40ac-bf16-b911f6d7b39f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "spark-dl-tf",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
